{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9c48213d-6e6a-4c10-838a-2a7c710c3a05",
   "metadata": {},
   "source": [
    "# Multi-Step Query Engine\n",
    "\n",
    "We have a multi-step query engine that's able to decompose a complex query into sequential subquestions. This\n",
    "guide walks you through how to set it up!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "50d3b817-b70e-4667-be4f-d3a0fe4bd119",
   "metadata": {},
   "source": [
    "#### Load documents, build the VectorStoreIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "690a6918-7c75-4f95-9ccc-d2c4a1fe00d7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    LLMPredictor,\n",
    "    ServiceContext,\n",
    ")\n",
    "from llama_index.llms import OpenAI\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48da73f-aadb-480c-8db1-99c915b7cc1c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# LLM Predictor (gpt-3)\n",
    "gpt3 = OpenAI(temperature=0, model=\"text-davinci-003\")\n",
    "service_context_gpt3 = ServiceContext.from_defaults(llm=gpt3)\n",
    "\n",
    "# LLMPredictor (gpt-4)\n",
    "gpt4 = OpenAI(temperature=0, model=\"gpt-4\")\n",
    "service_context_gpt4 = ServiceContext.from_defaults(llm=gpt4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03d1691e-544b-454f-825b-5ee12f7faa8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"../paul_graham_essay/data\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad144ee7-96da-4dd6-be00-fd6cf0c78e58",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b6caf93b-6345-4c65-a346-a95b0f1746c4",
   "metadata": {},
   "source": [
    "#### Query Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95d989ba-0c1d-43b6-a1d3-0ea7135f43a6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.query.query_transform.base import StepDecomposeQueryTransform\n",
    "\n",
    "# gpt-4\n",
    "step_decompose_transform = StepDecomposeQueryTransform(llm_predictor_gpt4, verbose=True)\n",
    "\n",
    "# gpt-3\n",
    "step_decompose_transform_gpt3 = StepDecomposeQueryTransform(\n",
    "    llm_predictor_gpt3, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a124db0-e2d7-4566-bcec-1d41cf669ff4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "index_summary = \"Used to answer questions about the author\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85466fdf-93f3-4cb1-a5f9-0056a8245a6f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "from llama_index.query_engine.multistep_query_engine import MultiStepQueryEngine\n",
    "\n",
    "query_engine = index.as_query_engine(service_context=service_context_gpt4)\n",
    "query_engine = MultiStepQueryEngine(\n",
    "    query_engine=query_engine,\n",
    "    query_transform=step_decompose_transform,\n",
    "    index_summary=index_summary,\n",
    ")\n",
    "response_gpt4 = query_engine.query(\n",
    "    \"Who was in the first batch of the accelerator program the author started?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdda1b2c-ae46-47cf-91d7-3153e8d0473b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response_gpt4}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9670bd-729d-478b-a77c-c6e13c282456",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sub_qa = response_gpt4.metadata[\"sub_qa\"]\n",
    "tuples = [(t[0], t[1].response) for t in sub_qa]\n",
    "print(tuples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec88df57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response_gpt4 = query_engine.query(\n",
    "    \"In which city did the author found his first company, Viaweb?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "653508f1-b2b0-479a-85b3-113cda507231",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(response_gpt4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa93cdb-7007-4664-853a-5c81c6c17560",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(service_context=service_context_gpt3)\n",
    "query_engine = MultiStepQueryEngine(\n",
    "    query_engine=query_engine,\n",
    "    query_transform=step_decompose_transform_gpt3,\n",
    "    index_summary=index_summary,\n",
    ")\n",
    "\n",
    "response_gpt3 = query_engine.query(\n",
    "    \"In which city did the author found his first company, Viaweb?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05899fcf-7a04-4d21-9e6d-04983755d175",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(response_gpt3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a43c659-ffd0-40df-b52b-032e6647cf9f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
